package com.atguigu.sql.udf

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, StructField, StructType}

/**
 * description ：求平均值自定义uef函数
 * author      ：剧情再美终是戏 
 * mail        : 13286520398@163.com
 * date        ：Created in 2020/1/11 11:40
 * modified By ：
 * version:    : 1.0
 */
object MyAvg extends UserDefinedAggregateFunction {

  // 输入参数类型定义
  override def inputSchema: StructType = StructType(StructField("input", DoubleType) :: Nil)

  // 缓存区类型定义
  override def bufferSchema: StructType = StructType(StructField("sum", DoubleType) :: StructField("count", DoubleType) :: Nil)

  // 返回数据的类型定义
  override def dataType: DataType = DoubleType

  // 如果此函数是确定性的，则返回true，即给定相同的输入，*始终返回相同的输出
  override def deterministic: Boolean = false

  // 初始化
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0d
    buffer(1) = 0d
  }

  // 分区内聚合逻辑
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    buffer(0) = buffer.getDouble(0) + input.getDouble(0)
    buffer(1) = buffer.getDouble(1) + 1d
  }

  // 分区间聚合逻辑
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getDouble(0) + buffer2.getDouble(0)
    buffer1(1) = buffer1.getDouble(1) + buffer2.getDouble(1)
  }

  // 最终的返回结果
  override def evaluate(buffer: Row): Double = buffer.getDouble(0) / buffer.getDouble(1)
}
